from mindquantum.core.gates import H, X
from mindquantum.core.circuit import Circuit
from mindquantum.simulator import Simulator
from math import sqrt

def create_bell_statevector():
    """ Return a phi+ Bell statevector.
    """
    circuit = Circuit()
    circuit += H.on(0)
    circuit += X.on(1, 0)
    simulator = Simulator('mqvector', 2)
    simulator.apply_circuit(circuit)
    return simulator.get_qs()

def check(candidate):
    result = candidate()
    solution = create_bell_statevector()
    assert all(abs(r - s) < 1e-10 for r, s in zip(result, solution))

    print("OK")

if __name__ == '__main__':
    check(create_bell_statevector)